#include <stdio.h>
#include <string.h>
#include <stdint.h>
#include "api.h"
#include "rand.h"
#include "bin-lwe.h"
#include "ntt.h"
#include "ntt-1024.h"
#include "fips202.h"

#define CTESTS 100000
uint64_t loop=1;

static void print_uint64(unsigned long long num)
{
	if(num>=10)
		print_uint64(num/10);
	printf("%u",(unsigned int)(num%10));
}


long long int error_bit_num(unsigned char *k1, unsigned char *k2, int num)
{
	int i;
	long long sum=0;
	unsigned char temp;
	for(i=0;i<num;i++)
	{
		temp=k1[i]^k2[i];
		if(temp>0)
		{
			sum+=(temp&0x1);
			sum+=((temp>>1)&0x1);
			sum+=((temp>>2)&0x1);
			sum+=((temp>>3)&0x1);
			sum+=((temp>>4)&0x1);
			sum+=((temp>>5)&0x1);
			sum+=((temp>>6)&0x1);
			sum+=((temp>>7)&0x1);
		}
	}
	
	return sum;
}
int print_error_bit(unsigned char *k1, unsigned char *k2, int num)
{
	int i,sum=0;
	unsigned char temp;
	printf("\nerror bit:\n");
	for(i=0;i<num;i++)
	{
		temp=k1[i]^k2[i];
		printf("%d%d%d%d%d%d%d%d",temp&0x1,((temp>>1)&0x1),((temp>>2)&0x1),((temp>>3)&0x1),((temp>>4)&0x1),((temp>>5)&0x1),((temp>>6)&0x1),((temp>>7)&0x1));
	}
	printf("\n");
	
	return sum;
}

//print bytes
int print_bytes(unsigned char *buf, int len)
{
	int i;
	for(i=0;i<len;i++)
	{
		printf("%d ",buf[i]);
	}
	printf("\n");
	
	return 0;
}


//test correctness of pke
int test_pke_correctness()
{
	printf("\n =====================================  test pke  ========================================== \n");

	unsigned char pk[CRYPTO_PUBLICKEYBYTES];
	unsigned char sk[CRYPTO_SECRETKEYBYTES];
	unsigned char k1[CRYPTO_BYTES], k2[CRYPTO_BYTES], c[CRYPTO_CIPHERTEXTBYTES];
	int i, j;
	long long int  error_bit, sum = 0;
	long long int  error_num = 0, sum_bits;
	unsigned long long mlen=MESSAGE_LEN,clen=CRYPTO_CIPHERTEXTBYTES;

	printf("correctness test of pke:\n");
	for (j = 0; j < loop; j++)
	{
		crypto_encrypt_keypair(pk,sk);
		random_bytes(k1, CRYPTO_BYTES);
		for (i = 0; i < CTESTS; i++)
		{
			unsigned char seed[SEED_LEN];
			//generate seed
			random_bytes(seed,SEED_LEN);
			crypto_encrypt(c, &clen, k1, mlen, pk, seed);
			// printf("c_byts=%lld\n",c_byts);
			crypto_encrypt_open(k2, &mlen, c, clen, sk);
			
			if (memcmp(k1, k2, mlen) != 0)
			{
				error_num++;
				error_bit = error_bit_num(k1, k2, mlen);
				sum += error_bit;
				if (error_bit > 0)
				{
					printf("error bit num:");
					print_uint64(error_bit);
					printf("\n");
					print_error_bit(k1, k2, mlen);
				}
			}
		}
		printf("test %d error block:", j + 1);
		print_uint64(error_num);
		printf(" error bit:");
		print_uint64(sum);
		printf("\n");

	}
	sum_bits = CTESTS * loop*CRYPTO_BYTES * 8;
	printf("total error bit:");
	print_uint64(sum);
	printf("/");
	print_uint64(sum_bits);
	printf("\n\n");

	// printf("pke test over\n\n");
	return error_num;
}


//test kem fo correctness
int test_kem_fo_correctness()
{
	unsigned char pk[CRYPTO_PUBLICKEYBYTES];
	unsigned char sk[CRYPTO_SECRETKEYBYTES];
	unsigned char k1[CRYPTO_SESSION],k2[CRYPTO_SESSION],c[CRYPTO_CIPHERTEXTBYTES];
	int i,j;
	long long int  error_num=0;
	
	printf("correctness test of kem_fo:\n");
	for(j=0;j<loop;j++)
	{
		crypto_kem_keypair(pk,sk);
		random_bytes(k1,CRYPTO_SESSION);
		for(i=0;i<CTESTS;i++)
		{
			crypto_kem_enc(c,k1,pk);
			crypto_kem_dec(k2,c,sk);
			
			if(memcmp(k1,k2,CRYPTO_SESSION)!=0)
			{
				error_num++;
			}
			
		}
		printf("test %d error block:",j+1);
		print_uint64(error_num);
		printf("\n");
	}
	printf("\n");

	return error_num;
}


int test_mul_correctness()
{
	//   test correctness of polymul
	unsigned char ctestn[DIM_N];
	uint16_t ctest[DIM_N];
	// const unsigned char uinta[DIM_N]={121,4,42,51,250,65,26,55,151,58,236,160,222,107,43,22,164,95,161,227,245,76,0,132,40,131,235,96,206,68,233,156,76,162,110,112,204,38,104,248,153,90,161,159,222,153,185,128,11,182,131,20,157,114,147,244,209,73,243,240,158,73,128,118,132,1,30,121,158,112,112,97,124,19,206,45,100,202,8,183,193,213,97,241,149,93,27,146,140,210,112,239,38,36,1,216,101,99,104,113,9,78,123,1,54,25,133,19,22,197,12,230,113,66,191,129,0,209,89,40,161,4,233,141,173,43,181,98,141,15,164,83,68,25,4,226,202,9,38,231,100,90,171,74,187,105,59,233,156,100,162,102,148,21,121,105,143,161,8,150,220,151,151,231,175,41,244,74,135,234,192,60,221,47,7,155,184,79,18,0,114,26,162,49,101,148,206,34,178,69,101,19,91,229,152,88,142,228,158,195,91,214,37,139,215,193,179,220,3,194,228,81,99,2,66,159,220,66,191,75,15,221,245,130,187,193,113,234,160,104,189,220,35,180,110,5,253,46,200,26,230,183,118,157,118,60,117,197,142,18,160,199,191,26,40,120,181,225,78,13,243,65,68,232,178,74,143,208,13,46,123,129,0,11,115,124,193,58,189,174,22,185,55,43,133,24,191,26,73,48,235,101,161,17,49,32,135,246,85,118,51,212,11,145,164,158,229,4,158,19,173,153,247,9,67,162,244,8,97,52,225,24,223,129,204,207,180,242,86,175,43,102,55,55,56,114,44,123,249,119,56,161,36,67,33,219,61,209,109,93,100,113,244,89,145,17,102,39,154,152,246,132,43,93,185,234,125,167,179,234,192,176,165,166,120,55,194,180,2,71,8,211,137,5,78,128,190,109,95,37,98,153,212,33,10,71,223,29,118,111,67,53,32,225,174,38,118,118,71,83,80,171,29,242,136,189,34,210,178,116,227,234,159,51,248,87,187,18,255,61,99,206,163,88,148,104,184,19,125,27,70,86,54,130,115,94,145,178,193,120,225,147,121,94,69,74,34,118,143,176,67,11,225,106,133,243,151,26,40,113,19,213,81,79,116,80,112,117,123,8,136,186,101,179,139,144,173,190,162,147,241,30,124,227,0,220,84,41,192,13,246,181,45,124,35,37,243,151,75,250,145,103};
	// const unsigned char uints[DIM_N]={250,0,250,1,0,0,0,0,0,0,250,0,0,0,250,250,0,0,0,250,1,0,0,0,0,0,250,250,0,0,0,0,250,1,0,1,250,0,1,0,1,0,1,0,0,250,0,1,1,0,0,0,0,1,1,0,250,250,0,1,0,1,1,250,0,1,1,0,1,250,1,0,0,250,0,1,1,1,1,0,0,0,250,250,0,0,1,0,1,0,0,1,250,0,1,250,0,0,1,1,250,0,0,0,0,1,1,250,250,0,1,1,0,0,1,250,0,0,250,0,250,1,250,250,1,250,0,1,1,1,250,1,0,1,250,0,0,0,250,0,0,250,250,250,250,250,250,1,0,1,0,0,250,0,0,250,250,1,1,1,1,1,0,0,1,0,1,0,250,1,0,250,0,0,0,0,0,1,0,0,250,250,250,250,0,250,0,250,0,0,0,250,0,250,250,1,0,0,250,250,1,0,0,0,0,0,1,250,250,0,1,0,0,250,250,1,0,1,1,1,250,250,1,0,0,0,0,250,0,0,250,0,0,250,0,250,250,1,0,0,0,0,250,0,0,0,0,0,0,250,250,0,1,0,250,0,0,250,250,250,250,250,1,250,250,0,250,250,1,0,250,250,250,0,250,0,0,0,250,0,1,0,0,250,1,0,0,1,1,1,0,0,250,250,0,0,0,0,250,0,0,250,0,0,0,250,1,1,0,0,0,250,1,250,1,1,0,250,250,1,250,1,0,0,0,250,0,1,1,0,1,1,250,1,0,0,0,0,0,1,0,0,0,1,0,250,1,1,0,0,1,0,0,1,250,1,0,250,0,0,1,0,1,1,0,250,0,1,1,0,1,0,0,1,0,0,0,0,0,250,1,250,250,250,0,0,250,1,0,0,0,1,250,0,0,250,0,0,0,250,0,1,250,250,0,250,0,0,1,1,0,0,250,0,1,250,0,0,250,1,0,250,0,1,1,0,250,250,1,0,1,0,250,250,0,1,250,250,1,1,250,0,250,250,250,250,250,1,0,250,1,0,0,250,250,250,1,250,0,0,1,1,0,250,1,0,1,0,1,0,250,250,0,1,250,0,0,1,1,0,0,0,0,1,250,0,0,0,250,1,1,0,1,250,0,1,0,0,0,0,0,0,0,0,1,250,0,0,0,0,0,250};
	// uint16_t uinta2[DIM_N]={121,4,42,51,250,65,26,55,151,58,236,160,222,107,43,22,164,95,161,227,245,76,0,132,40,131,235,96,206,68,233,156,76,162,110,112,204,38,104,248,153,90,161,159,222,153,185,128,11,182,131,20,157,114,147,244,209,73,243,240,158,73,128,118,132,1,30,121,158,112,112,97,124,19,206,45,100,202,8,183,193,213,97,241,149,93,27,146,140,210,112,239,38,36,1,216,101,99,104,113,9,78,123,1,54,25,133,19,22,197,12,230,113,66,191,129,0,209,89,40,161,4,233,141,173,43,181,98,141,15,164,83,68,25,4,226,202,9,38,231,100,90,171,74,187,105,59,233,156,100,162,102,148,21,121,105,143,161,8,150,220,151,151,231,175,41,244,74,135,234,192,60,221,47,7,155,184,79,18,0,114,26,162,49,101,148,206,34,178,69,101,19,91,229,152,88,142,228,158,195,91,214,37,139,215,193,179,220,3,194,228,81,99,2,66,159,220,66,191,75,15,221,245,130,187,193,113,234,160,104,189,220,35,180,110,5,253,46,200,26,230,183,118,157,118,60,117,197,142,18,160,199,191,26,40,120,181,225,78,13,243,65,68,232,178,74,143,208,13,46,123,129,0,11,115,124,193,58,189,174,22,185,55,43,133,24,191,26,73,48,235,101,161,17,49,32,135,246,85,118,51,212,11,145,164,158,229,4,158,19,173,153,247,9,67,162,244,8,97,52,225,24,223,129,204,207,180,242,86,175,43,102,55,55,56,114,44,123,249,119,56,161,36,67,33,219,61,209,109,93,100,113,244,89,145,17,102,39,154,152,246,132,43,93,185,234,125,167,179,234,192,176,165,166,120,55,194,180,2,71,8,211,137,5,78,128,190,109,95,37,98,153,212,33,10,71,223,29,118,111,67,53,32,225,174,38,118,118,71,83,80,171,29,242,136,189,34,210,178,116,227,234,159,51,248,87,187,18,255,61,99,206,163,88,148,104,184,19,125,27,70,86,54,130,115,94,145,178,193,120,225,147,121,94,69,74,34,118,143,176,67,11,225,106,133,243,151,26,40,113,19,213,81,79,116,80,112,117,123,8,136,186,101,179,139,144,173,190,162,147,241,30,124,227,0,220,84,41,192,13,246,181,45,124,35,37,243,151,75,250,145,103};
	// uint16_t uints2[DIM_N]={250,0,250,1,0,0,0,0,0,0,250,0,0,0,250,250,0,0,0,250,1,0,0,0,0,0,250,250,0,0,0,0,250,1,0,1,250,0,1,0,1,0,1,0,0,250,0,1,1,0,0,0,0,1,1,0,250,250,0,1,0,1,1,250,0,1,1,0,1,250,1,0,0,250,0,1,1,1,1,0,0,0,250,250,0,0,1,0,1,0,0,1,250,0,1,250,0,0,1,1,250,0,0,0,0,1,1,250,250,0,1,1,0,0,1,250,0,0,250,0,250,1,250,250,1,250,0,1,1,1,250,1,0,1,250,0,0,0,250,0,0,250,250,250,250,250,250,1,0,1,0,0,250,0,0,250,250,1,1,1,1,1,0,0,1,0,1,0,250,1,0,250,0,0,0,0,0,1,0,0,250,250,250,250,0,250,0,250,0,0,0,250,0,250,250,1,0,0,250,250,1,0,0,0,0,0,1,250,250,0,1,0,0,250,250,1,0,1,1,1,250,250,1,0,0,0,0,250,0,0,250,0,0,250,0,250,250,1,0,0,0,0,250,0,0,0,0,0,0,250,250,0,1,0,250,0,0,250,250,250,250,250,1,250,250,0,250,250,1,0,250,250,250,0,250,0,0,0,250,0,1,0,0,250,1,0,0,1,1,1,0,0,250,250,0,0,0,0,250,0,0,250,0,0,0,250,1,1,0,0,0,250,1,250,1,1,0,250,250,1,250,1,0,0,0,250,0,1,1,0,1,1,250,1,0,0,0,0,0,1,0,0,0,1,0,250,1,1,0,0,1,0,0,1,250,1,0,250,0,0,1,0,1,1,0,250,0,1,1,0,1,0,0,1,0,0,0,0,0,250,1,250,250,250,0,0,250,1,0,0,0,1,250,0,0,250,0,0,0,250,0,1,250,250,0,250,0,0,1,1,0,0,250,0,1,250,0,0,250,1,0,250,0,1,1,0,250,250,1,0,1,0,250,250,0,1,250,250,1,1,250,0,250,250,250,250,250,1,0,250,1,0,0,250,250,250,1,250,0,0,1,1,0,250,1,0,1,0,1,0,250,250,0,1,250,0,0,1,1,0,0,0,0,1,250,0,0,0,250,1,1,0,1,250,0,1,0,0,0,0,0,0,0,0,1,250,0,0,0,0,0,250};
	
	unsigned char uinta[DIM_N];
	unsigned char uints[DIM_N];
	uint16_t uinta2[DIM_N];
	uint16_t uints2[DIM_N];
	// const unsigned char *pk;
	unsigned char seed[SEED_LEN];
	random_bytes(seed,SEED_LEN);
	unsigned char seeds[3*SEED_LEN];
	keccak_state state;
	shake256_absorb_once(&state, seed, SEED_LEN);
	shake256_squeeze(seeds, 3*SEED_LEN, &state);
	gen_a(uinta,seeds);
    shake256_absorb_once(&state, seeds + SEED_LEN, SEED_LEN);
	gen_e(uints,&state);

	// unsigned char seed[SEED_LEN];
	// random_bytes(seed,SEED_LEN);
	// gen_a(uinta,seed);
	// gen_e(uints,seed);
	int i;
	for(i=0;i<DIM_N;i++)
	{
		uinta2[i]=uinta[i];
		uints2[i]=uints[i];
	}
	// printf(" input a: \n"); //
	// for (i = 0; i < DIM_N; i++)
	// {

	// 	printf("%d,", uinta[i]);
	// }
	// printf("\n");
	// printf(" input r: \n"); //
	// int num=0;
	// for (i = 0; i < DIM_N; i++)
	// {
	// 	if(uints[i]==0)
	// 		num++;
	// 	printf("%d,", uints[i]);
	// }
	// printf("\n number of 0 : %d \n",num);

	printf("\n =====================================  test poly_mul======================================= \n");
	poly_mul(uinta,uints,ctestn,DIM_N);
	NormalMul_unsign_251(uinta2,uints2,ctest);

	printf("\n");
	
	// for(i=0;i<DIM_N;i++)
	// 	ctest[i]=ctest[i]+e1[i];
	
	for(i=0;i<DIM_N;i++)
	{
		if(ctestn[i]!=ctest[i])
			break;
	}
	if(i<DIM_N)
	{
		printf("Lifting Multiplication            : wrong , wrong index is %d, the wrong value is %d, correct valure is %d\n", i, ctest[i], ctestn[i]);
	}
	else
	{
		printf("Lifting Multiplication            : right \n");
	}


	// printf("\n ntt res: \n"); //
	// for (i = 0; i < DIM_N; i++)
	// {

	// 	printf("%d,", ctestn[i]);
	// }
	// printf("\n mul res: \n"); //
	// for (i = 0; i < DIM_N; i++)
	// {

	// 	printf("%d,", ctest[i]);
	// }

	/*
	unsigned char ctesto[DIM_N];
	poly_mul_original(uinta,uints,ctesto,DIM_N);
	for(i=0;i<DIM_N;i++)
	{
		if(ctestn[i]!=ctesto[i])
			break;
	}
	if(i<DIM_N)
	{
		printf("Lifting Multiplication            : wrong , wrong index is %d, the wrong value is %d, correct valure is %d\n", i, ctestn[i], ctesto[i]);
	}
	else
	{
		printf("Lifting Multiplication            : right \n");
	}
	printf("\n original res: \n"); //
	for (i = 0; i < DIM_N; i++)
	{

		printf("%d,", ctesto[i]);
	}
	*/



	printf("\n =====================================  test poly_aff======================================= \n");
	unsigned char e[DIM_N];
	gen_e(e,&state);
	// printf(" input e: \n"); //
	// for (i = 0; i < DIM_N; i++)
	// {

	// 	printf("%d,", e[i]);
	// }

	for(i=0;i<DIM_N;i++)
	{
		if(e[i]>q_half)
			e[i] = 250;
		else
			e[i]=e[i];
	}
	poly_aff(uinta, uints, e, ctestn, DIM_N);

	for(i=0;i<DIM_N;i++)
		ctest[i]=(ctest[i]+e[i] )% Q;
	
	for(i=0;i<DIM_N;i++)
	{
		if(ctestn[i]!=ctest[i])
			break;
	}
	if(i<DIM_N)
	{
		printf("Lifting Multiplication aff            : wrong , wrong index is %d, the wrong value is %d, correct valure is %d\n", i, ctest[i], ctestn[i]);
	}
	else
	{
		printf("Lifting Multiplication aff            : right \n");
	}

	/*
	poly_aff_original(uinta, uints, e, ctesto, DIM_N);
	for(i=0;i<DIM_N;i++)
	{
		if(ctestn[i]!=ctesto[i])
			break;
	}
	if(i<DIM_N)
	{
		printf("Lifting Multiplication aff            : wrong , wrong index is %d, the wrong value is %d, correct valure is %d\n", i, ctestn[i], ctesto[i]);
	}
	else
	{
		printf("Lifting Multiplication aff            : right \n");
	}

	printf("\n original res: \n"); //
	for (i = 0; i < DIM_N; i++)
	{

		printf("%d,", ctesto[i]);
	}
	printf("\n ntt res: \n"); //
	for (i = 0; i < DIM_N; i++)
	{

		printf("%d,", ctestn[i]);
	}
	printf("\n normal res: \n"); //
	for (i = 0; i < DIM_N; i++)
	{

		printf("%d,", ctest[i]);
	}
	*/

	return 0;
}